package com.study.feature

import org.apache.spark.ml.feature.{RegexTokenizer, StopWordsRemover, StringIndexer, Tokenizer}
import org.apache.spark.sql.SparkSession

/**
 * 文本特征处理
 * 1、分词
 * 2、去掉停用词
 * 3、词稀疏编码
 *
 * @author stephen
 * @date 2019-08-26 09:59
 */
object TextFeatureDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("TextFeatureDemo")
      .master("local[*]")
      .getOrCreate()

    // 定义数据源
    val sentenceDF = spark.createDataFrame(Seq(
      (0, "You must give any other recipients of the Work or Derivative Works a copy of this License"),
      (1, "You must cause any modified files to carry prominent notices stating that You changed the files"),
      (2, "You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works")
    )).toDF("label", "sentence")

    // 定义分词器
    // 默认配置的分词器
    val tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words")
    // 基于正则表达式的分词器
    val regexTokenizer = new RegexTokenizer().setInputCol("sentence").setOutputCol("words").setPattern("\\W").setGaps(false)

    // 进行分词
    val tokenized = tokenizer.transform(sentenceDF)
    tokenized.select("label", "words").take(5).foreach(println)
    regexTokenizer.transform(sentenceDF).select("label", "words").take(5).foreach(println)

    // 定义停用词
    val stopWordsRemover = new StopWordsRemover().setInputCol("words").setOutputCol("filtered")
    // 去掉停用词
    val filtered = stopWordsRemover.transform(tokenized)
    filtered.select("label","filtered").show(false)

    // 词稀疏编码
    val df = spark.createDataFrame(Seq((1,"a"),(2,"b"),(3,"c"),(4,"a"),(5,"b"),(6,"d"),(7,"a"),(8,"f"),(9,"c"))).toDF("id","word")
    val indexer = new StringIndexer().setInputCol("word").setOutputCol("wordIndex")
    val indexed = indexer.fit(df).transform(df)
    indexed.select("id","word","wordIndex").show(false)
  }
}
